

# !/usr/bin/env python
# coding:utf-8
"""
@version:  1.0.0
@author:  Happylin
@file: re_fetch.py
@time: 2024/5/15 9:25
"""
import os
import re
import pandas as pd
import pdfplumber
from dateutil.parser import parse

class InvoiceExtract():



    def extract_invoice_date(self, text):
        # Identify potential date-related keywords
        keywords = ["开票日期"]
        for keyword in keywords:
            # print(keyword, text)
            pattern = rf"{keyword}\s*[：:]\s*(.*?)\n"
            matches = re.findall(pattern, text, re.DOTALL | re.IGNORECASE)
            if matches:
                for match in matches:
                    print("match:", match)
                    try:
                        # date = parse(match)  # Attempt to parse the extracted date
                        standard_date = match.replace("年", "-").replace("月", "-").replace("日", "")
                        return standard_date
                    except ValueError:
                        continue  # Ignore invalid date formats
        return None
    def extract_invoice_detail(self, pdf_path):
        invoice_header = {
            "电子发票文件名": os.path.basename(pdf_path),
            "发票类型": None,
            "发票代码": None,
            "发票号码": None,
            "开票日期": None,
            "购买方名称": None,
            "购买方-统一社会信用代码": None,
            "销售方名称": None,
            "销售方-统一社会信用代码": None,
            "备注": None,
            "含税金额": None,
            "总税额": None,
            "不含税金额": None,
            "发票备注": None,
            "税务局代码": None,
            "税务局名称": None
        }
        invoice_lines =  []

        invoice_line =   {
            "电子发票文件名": os.path.basename(pdf_path),
            "项目名称": None,
            "规格型号": None,
            "项目单位": None,
            "项目数量": None,
            "项目单价": None,
            "项目金额": None,
            "项目税率": None,
            "项目税额": None,
            "项目含税金额": None
        }
        with pdfplumber.open(pdf_path) as pdf:
            for page in pdf.pages:
                # 提取文本内容
                text = page.extract_text()
                # 判断发票类型，发票类型有'电子发票（增值税专用发票）'和'电子发票（普通发票）'以及其他地方的发票类型
                # 电子发票（增值税专用发票）的特点是：在text中找到包含'增值税专用发票'的字符串
                # 电子发票（普通发票）的特点是：在text中找到包含'普通发票'的字符串
                # 其他地方的发票类型的特点是：在text中找到包含'电子普通发票'的字符串和各个地方的地方名称的字符串
                # 设置一个变量invoice_type，用于存储发票类型
                invoice_type = None
                if '电子发票（增值税专用发票）' in text:
                    invoice_type = '电子发票（增值税专用发票）'
                elif '电子发票（普通发票）' in text:
                    invoice_type = '电子发票（普通发票）'
                elif '电子普通发票' in text:
                    invoice_type = '电子普通发票'
                else:
                    invoice_type = '其他发票'
                invoice_header["发票类型"] = invoice_type
                # print("invoice_type:", invoice_type)


                # print("text:", text)
                """
                如果text为空，说明无法提取文本内容，将'发票备注'字段写入'无法提取文本内容'，然后继续下一个循环
                发票代码、发票号码、开票日期、总金额、总税额、服务名称等字段的提取方法是：
                在text中找到包含'发票代码'、'发票号码'、'开票日期'、'合计'、'小写'等关键字的字符串，然后在这个字符串之后找到数字字符，然后提取这个数字字符
                这些对增值税发票和普通发票都适用
                但是发票行的提取会有所不同，因为增值税发票和普通发票的格式不同，所以提取方法也不同
                """
                if not text or  invoice_type == '其他发票':
                    invoice_header["发票备注"] = "无法提取文本内容"
                    continue
                invoice_header["发票备注"] = "成功提取文本内容"
                # print("text:", type(text), text)
                # 提取发票代码
                if not invoice_header["发票代码"]:
                    code_match = re.search(r'发票代码[:：]?\s*(\d+)', text)
                    if code_match:
                        invoice_header["发票代码"] = code_match.group(1).strip()
                # 提取发票号码
                if not invoice_header["发票号码"]:
                    number_match = re.search(r'发票号码[:：]?\s*(\d+)', text)
                    if number_match:
                        invoice_header["发票号码"] = number_match.group(1).strip()

                # 提取开票日期
                if not invoice_header["开票日期"]:
                    invoice_header["开票日期"] = self.extract_invoice_date(text)

                # 提取总金额    
                # 提取方法是：在text中找到包含'小写'的字符串，然后在这个字符串之后找到第一个'¥'字符，然后在这个字符之后找到数字字符，然后提取这个数字字符
                if not invoice_header["含税金额"]:
                    total_amount_match = re.search(r'小写.*?¥\s*(\d+\.\d+)', text)
                    if total_amount_match:
                        invoice_header["含税金额"] = total_amount_match.group(1).strip()


                # 提取总税额
                # 提取方法是：在text中找到包含'合 计'的字符串，然后在这个字符串之后找到第二个'¥'字符，然后在这个字符之后找到数字字符，然后提取这个数字字符
                if not invoice_header["总税额"]:
                    total_tax_match = re.search(r'合\s*计.*¥\s*\d+\.\d+.*¥\s*(\d+\.\d+)', text, re.DOTALL)
                    if total_tax_match:
                        invoice_header["总税额"] = total_tax_match.group(1).strip()
                # 提取总金额
                # 提取方法是：在text中找到包含'小写'的字符串，然后在这个字符串之后找到第一个''字符，然后在这个字符之后找到数字字符，然后提取这个数字字符
                if not invoice_header["不含税金额"]:
                    total_with_tax_match = re.search(r'小\s*写.¥\s*(\d+\.\d+)', text)
                    if total_with_tax_match:
                        invoice_header["不含税金额"] = total_with_tax_match.group(1).strip()

                # 判断发票类型

                # 提取表格数据
                tables = page.extract_tables()
                columns = []
                if tables:
                    for row in tables[0]:
                        for cell in row:
                            # print("cell:", cell)
                            cells = []
                            if cell != None:
                                for item in cell.split('\n'):
                                    cells.append(item)
                            
                            columns.append(cells)       

                    # print("columns:",columns)
                    """
                    将columns的元素是列表，其中一个列表包含['购', '买', '方']， 
                    此列表之后的列表元素中是的第一个原始是一个包含'名称'和冒号的字符串，
                    将此字符串用冒号作为分隔符，冒号后面的字符串 作为'购买方名称'写到invoice_header中的'购买方名称'字段
                    """
                    """
                    在现在的程序逻辑基础上增加一个逻辑，逻辑为：与‘名称’同一个列表中的下一个元素的值是包含'统一社会信用代码‘的值，
                    这个值用冒号分隔符分割后，第二个值作为’购买方名称-统一社会信用代码写入到invoice_header中
                    """
                    # 在下面程序逻辑上再增加如下逻辑：在columns中找到包含'销'、'售'、'方'的列表元素，
                    # 然后在这个列表元素之后的元素中找到包含'名'和'称'两个字符的元素，这个元素用冒号分隔符分割后，
                    # 第二个值作为'销售方名称'写入到invoice_header中的'销售方名称'字段，此元素的下一个元素中包含'统一社会信用代码'的值，
                    # 这个值用冒号分隔符分割后，第二个值作为'销售方名称-统一社会信用代码'写入到invoice_header中的'销售方名称-统一社会信用代码'字段
                    # 在上面程序逻辑上再增加如下逻辑：在columns中找到包含'货物或'字符的列表数据，刨除此列表中第一个元素，
                    # 后面的元素写入到invoice_line的'项目名称'字段。然后继续在columns中找到包含'规格型号'字符的列表数据，
                    # 刨除此列表中第一个元素，后面的元素写入到invoice_line的'规格型号'字段。
                    # 然后继续在columns中找到包含'单位'字符的列表数据，刨除此列表中第一个元素，
                    # 后面的元素写入到invoice_line的'单位'字段。然后继续在columns中找到包含'数 量'字符的列表数据，
                    # 刨除此列表中第一个元素，后面的元素写入到invoice_line的'数量'字段。
                    # 然后继续在columns中找到包含'单 价'字符的列表数据，刨除此列表中第一个元素，
                    # 后面的元素写入到invoice_line的'单价'字段。然后继续在columns中找到包含'金 额'字符的列表数据，
                    # 刨除此列表中第一个元素，后面的元素写入到invoice_line的'金额'字段。
                    # 然后继续在columns中找到包含'税率'字符的列表数据，刨除此列表中第一个元素，
                    # 后面的元素写入到invoice_line的'税率'字段。然后继续在columns中找到包含'税 额'字符的列表数据，
                    # 刨除此列表中第一个元素，后面的元素写入到invoice_line的'税额'字段。
                    for i, col in enumerate(columns):
                        # print("col:", col, '购' in col)
                        if col:
                            # 如果发票类型是'电子普通发票'，则按下面规则提取发票行数据
                            if invoice_type == '电子普通发票':
                                # print("电子普通发票col[0]:", col)
                                if '货物或应税劳务、服务名称' in col[0]:
                                    invoice_line["项目名称"] = col[1:-1]
                                elif '规格型号' in col[0]:
                                    invoice_line["规格型号"] = col[1:]
                                elif '单位' in col[0]:
                                    invoice_line["单位"] = col[1:]
                                elif '数 量' in col[0]:
                                    invoice_line["数量"] = col[1:]
                                elif '单 价' in col[0]:
                                    invoice_line["单价"] = col[1:]
                                elif '金 额' in col[0]:
                                    invoice_line["金额"] = col[1:-1]
                                elif '税率' in col[0]:
                                    invoice_line["税率"] = col[1:]
                                elif '税 额' in col[0]:
                                    invoice_line["税额"] = col[1:-1]
                               
                                
                                # print("电子普通发票invoice_lines:", invoice_lines)
                            # 如果发票类型是'电子发票（增值税专用发票）'或'电子发票（普通发票）'，则按下面规则提取发票行数据    
                            if invoice_type == '电子发票（增值税专用发票）' or invoice_type == '电子发票（普通发票）':
                                if '项目名称' in col[0]:
                                    # print("项目col:", col)
                                    #在col[0]列表中包含多个元素，第二个元素和最后一个元素之间是发票的项目名称，规格型号，单位，单价，数量，金额，税率，税额的信息，这些数据是用空格作为分隔符进行分割的，其中项目名称和单价，数量，金额，税率，税额是一定会有的数据，但是规格型号，单位不一定会有，所以在提取这些数据的时候，需要判断是否有这些数据，如果有这些数据，就提取这些数据，如果没有这些数据，就不提取这些数据
                                    #取col[0]列表中的第二个元素和最后一个元素之间的数据
                                    item_list = col[1:-1]
                                    for item in item_list:
                                        # print("item:", item)
                                        #用空格作为分隔符分割数据，并判断数据的长度
                                        item_parts = item.split()
                                        # print("item_parts:", item_parts)
                                        #如果数据的长度大于等于6，说明数据中包含项目名称，规格型号，单位，单价，数量，金额，税率，税额的信息
                                        if len(item_parts) >= 6:
                                            invoice_line["项目名称"] = item_parts[0]
                                            invoice_line["单价"] = item_parts[-5]
                                            invoice_line["数量"] = item_parts[-4]
                                            invoice_line["金额"] = item_parts[-3]
                                            invoice_line["税率"] = item_parts[-2]
                                            invoice_line["税额"] = item_parts[-1]
                                            if len(item_parts) > 6:
                                                invoice_line["规格型号"] = item_parts[1]
                                            if len(item_parts) > 7:
                                                invoice_line["单位"] = item_parts[2]
                                        # 增加发票行数据
                                        # print("电子发票（增值税专用发票）line:", invoice_line)
                                        # print("电子发票（增值税专用发票）line:", invoice_line)
                                        invoice_lines.append(invoice_line.copy())    
                                        # print("电子发票（增值税专用发票）invoice_lines:", invoice_lines)
                                        # print("电子发票（增值税专用发票）invoice_lines:", invoice_lines)

                    
                            # 提取购买方名称和销售方名称信息
                            if '购' in col and '买' in col and '方' in col:
                                for subsequent_col in columns[i+1:]:
                                    if any('名' in item and '称' in item for item in subsequent_col):
                                        buyer_info = subsequent_col[0].split('：', 1)
                                        if len(buyer_info) > 1:
                                            invoice_header["购买方名称"] = buyer_info[1].strip()
                                        # Check for the next element containing '统一社会信用代码'
                                        if len(subsequent_col) > 1:
                                            tax_id_info = subsequent_col[1].split('：', 1)
                                            if len(tax_id_info) > 1:
                                                invoice_header["购买方-统一社会信用代码"] = tax_id_info[1].strip()
                                        break                                
                            if '销' in col and '售' in col and '方' in col:
                                for subsequent_col in columns[i+1:]:
                                    if any('名' in item and '称' in item for item in subsequent_col):
                                        seller_info = subsequent_col[0].split('：', 1)
                                        if len(seller_info) > 1:
                                            invoice_header["销售方名称"] = seller_info[1].strip()
                                        # Check for the next element containing '统一社会信用代码'
                                        if len(subsequent_col) > 1:
                                            tax_id_info = subsequent_col[1].split('：', 1)
                                            if len(tax_id_info) > 1:
                                                invoice_header["销售方-统一社会信用代码"] = tax_id_info[1].strip()
                                        break   
                            # 提取备注信息，备注信息在                             
                            if '备' in col and '注' in col:
                                # print("subsequent_col:", columns[i+1:], columns[i+1:][0][0])
                                invoice_header["备注"] = columns[i+1:][0][0]

                    if invoice_type == '电子普通发票':
                        # invoice_line中的数据是一个字典，字典中的数据有的是列表，有的是字符串，有的是None，
                        # 需要根据列表的长度来判断是否需要将数据拆分成多行数据，然后将拆分后的数据写入到invoice_lines中，
                        # 其中如果数据是None或者是字符串的，
                        # 将none或者字符串重复多次写入到invoice_lines中
                        max_length = max(len(v) if isinstance(v, list) else 1 for v in invoice_line.values())
                        for i in range(max_length):
                            new_line = {}
                            for key, value in invoice_line.items():
                                if isinstance(value, list):
                                    new_line[key] = value[i] if i < len(value) else None
                                else:
                                    new_line[key] = value
                            invoice_lines.append(new_line)   
                        # print("invoice_lines:", invoice_lines)
                            

                        
                    # invoice_lines.append(invoice_line)
        # print("invoice_lines:", invoice_lines)                     
        return invoice_header, invoice_lines


def extract_from_directory(directory_path):
    invoices = {"Header": [], "Line": []}
   
    for filename in os.listdir(directory_path):
        if filename.endswith(".pdf"):
            invoice_headers = []
            invoice_line = []
            pdf_path = os.path.join(directory_path, filename)
            invoice_headers,invoice_line = InvoiceExtract().extract_invoice_detail(pdf_path)
            invoices["Header"].append(invoice_headers)
            invoices["Line"].append(invoice_line)
    #在返回数据前，将数据中的None值替换为''
    for invoice in invoices["Header"]:
        for key, value in invoice.items():
            if value == None:
                invoice[key] = ''
    for invoicelines in invoices["Line"]:
        for invoiceline in invoicelines:
            for key, value in invoiceline.items():
                if value == None:
                    invoiceline[key] = ''
    return invoices
  

if __name__ == '__main__':
    #print('PyCharm')
    print(extract_from_directory(r'D:\code\python\invoice_extract\pdf'))

